
import torch
from torch import nn
from model import makeResNet152, makeResNet50, makeNoClassiferResNet152, makeNoClassiferResNet50, CircleModel
from torch import optim
from torch.utils.data import DataLoader
from dataset import classiferDataset, VeRiClassiferDataset, VeRiCircleDataset
from loss import FocalLoss, CircleLoss, convert_label_to_similarity

def main():
    device = torch.device("cuda")
    maxAcc = 0
    maxLoss = 9999999999
    batchsz = 32
    classNum = 576

    trainDataPath = ".\\data\\VeRi\\image_train\\"
    testDataPath = ".\\data\\VeRi\\image_test\\"

    trainDataSet = VeRiCircleDataset(trainDataPath)
    print(len(trainDataSet))
    trainDataSet = DataLoader(trainDataSet, batch_size=batchsz, shuffle=True, num_workers=4)

    testDataSet = VeRiCircleDataset(testDataPath)
    print(len(testDataSet))
    testDataSet = DataLoader(testDataSet, batch_size=batchsz*8, shuffle=True, num_workers=4)

    # model = makeResNet152(classNum).to(device)
    # model = makeNoClassiferResNet152().to(device)
    # model = makeNoClassiferResNet50().to(device)
    # model = CircleModel().to(device)
    model = torch.load(".\\savedModels2\\model_epoch_50_right_0.9995235322145164.pkl")
    model = nn.Sequential(*list(model.children())[:-1],)
    print(model)
    p = sum(map(lambda p:p.numel(), model.parameters()))
    print('parameters size:', p)

    startEpoch = 51


    _lr = 0
    optimizer = optim.Adam(model.parameters(), lr=_lr , betas=(0.9, 0.999), eps=1e-09)
    # criterion = nn.CrossEntropyLoss().to(device)
    # criterion = FocalLoss(classNum).to(device)
    criterion = CircleLoss(m=0.25, gamma=256).to(device)

    for epoch in range(startEpoch, 10000):

        if epoch >= 0:
            _lr = 1e-4
        if epoch >= 100:
            _lr = 1e-5
        if epoch >= 150:
            _lr = 1e-6
        if epoch >= 200:
            _lr = 1e-7
        for p in optimizer.param_groups:
            p['lr'] = _lr


        trainTotalLoss = 0
        trainTotalStep = 0
        model.train()
        for batchidx, (x, label) in enumerate(trainDataSet):

            x, label = x.to(device), label.to(device)
            logits = model(x)
            label = label.long()

            # loss = criterion(logits, label)
            inp_sp, inp_sn = convert_label_to_similarity(logits, label)
            loss = criterion(inp_sp, inp_sn)
            # loss = criterion(logits, label)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            trainTotalLoss += float(loss)
            trainTotalStep += 1

            if batchidx % 100 == 0:
                # print(float(loss))
                print(trainTotalLoss / trainTotalStep)


        model.eval()
        testTotalLoss = 0
        testTotalStep = 0
        testTotalRight = 0
        testTotalNum = 0
        for batchidx, (x, label) in enumerate(testDataSet):
            with torch.no_grad():
                x, label = x.to(device), label.to(device)
                logits = model(x)

                test_inp_sp, test_inp_sn = convert_label_to_similarity(logits, label)
                test_loss = criterion(test_inp_sp, test_inp_sn)
                testTotalLoss += float(test_loss)
                testTotalStep += 1
                # logits = torch.softmax(logits, dim=-1)
                # logits = torch.argmax(logits, dim=-1)
                #
                # same = (logits == label).long().sum()
                # testTotalRight += same.item()
                # testTotalNum += logits.shape[0]

        # testAcc = testTotalRight / testTotalNum


        # print("epoch: ", epoch, "  testAcc: ", testAcc)
        # if epoch > 0 and testAcc >= maxAcc:
        #     torch.save(model, ".\\savedModels\\model_epoch_{}_right_{}.pkl".format(epoch, testAcc))
        #     maxAcc = testAcc
        # if epoch % 1 == 0:
        #     torch.save(model, ".\\savedModels\\n_model_epoch_{}_right_{}.pkl".format(epoch, testAcc))


        trainLoss = trainTotalLoss / trainTotalStep
        testLoss = testTotalLoss / testTotalStep
        print("epoch: ", epoch, "   train loss: ", trainLoss, "   test loss: ", testLoss)
        if epoch > 0 and  testLoss < maxLoss:
            torch.save(model, ".\\savedModels\\model_epoch_{}_trainLoss_{}_testLoss_{}.pkl".format(epoch, trainLoss, testLoss))
            maxLoss = testLoss
        if epoch % 1 == 0:
            torch.save(model, ".\\savedModels\\n_model_epoch_{}_trainLoss_{}_testLoss_{}.pkl".format(epoch, trainLoss, testLoss))



if __name__ == '__main__':
    main()
